from ultralytics import YOLO
from PIL import Image

# 加载预训练的YOLOv8n模型
model = YOLO('./results_runs/dataSet_good2/E10_IZ64/classify/train/weights/best.pt')

# 在图片上运行推理
results = model('./dataSet_good/val/battery/battery1.jpg')


# 这里Results只识别但张照片，我们用result[0]引用第一张的结果
if len(results) > 0 and results[0].probs is not None:
    top_class_id = results[0].probs.top1
    class_name = results[0].names[top_class_id]
    confidence = results[0].probs.data[top_class_id]
    im_array = results[0].plot()                        # 绘制包含预测结果的BGR numpy数组
    im = Image.fromarray(im_array[..., ::-1])           # RGB PIL图像
    im.show()                                           # 显示图像
    im.save('results3.jpg')                             # 保存图像
    print(f"识别结果: {class_name} (类别序号: {top_class_id}, 置信度: {confidence:.2f})")

